[Triton] unify the loading and add a UT for configs by Boss2002n · Pull Request #4630 · ROCm/aiter

Boss2002n · 2026-08-07T17:10:15Z

Motivation

Technical Details

Test Plan

Test Result

Submission Checklist

Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

Copilot

Pull request overview

This PR centralizes Triton JSON config loading/caching via a shared loader and updates multiple kernels/tests/benchmarks to use the unified path (plus a few correctness/robustness tweaks around config mutation and device capability gating).

Changes:

Introduces utils/core.py::load_config_json and migrates multiple config loaders to use it (with consistent caching and fallback logic).
Switches several GEMM call sites (tests/benchmarks and some kernels) to use get_gemm_config(...) instead of private per-kernel config loaders.
Fixes config mutation hazards by copying shared cached configs in attention MLA decode/rope paths.

Reviewed changes

Copilot reviewed 45 out of 45 changed files in this pull request and generated 4 comments.

Show a summary per file

File	Description
op_tests/triton_tests/gemm/basic/test_gemm_a8w8.py	Uses unified `get_gemm_config` in split-K tests.
op_tests/triton_tests/attention/test_mla_decode_rope.py	Deep-copies shared cached config to avoid test cross-talk.
op_tests/triton_tests/attention/test_la.py	Aligns test with updated `_get_config()` calling convention.
op_tests/op_benchmarks/triton/bench_moe_gemm_a8w8_blockscale.py	Uses unified `get_gemm_config` + split-K param computation.
aiter/ops/triton/utils/core.py	Adds shared cached JSON loader (`load_config_json`).
aiter/ops/triton/utils/gemm_config_utils.py	Refactors GEMM config resolution to use `load_config_json` and unified caching semantics.
aiter/ops/triton/utils/conv_config_utils.py	Switches conv config loading to `load_config_json`.
aiter/ops/triton/utils/moe_config_utils.py	Switches MoE config loading + caching policy integration (`USE_LRU_CACHE`).
aiter/ops/triton/utils/mhc_config_utils.py	Refactors MHC config loading with gfx942 fallback + shared JSON loader.
aiter/ops/triton/utils/_triton/gemm_tune_check.py	Fixes module-path mapping for locating kernel modules.
aiter/ops/triton/moe/moe_op_gemm_a8w4.py	Uses shared JSON loader for dispatch table.
aiter/ops/triton/gluon/gemm_afp4wfp4.py	Routes config loading through `get_gemm_config(..., backend="gluon")`.
aiter/ops/triton/gluon/gemm_a8w8.py	Routes config loading through `get_gemm_config(..., backend="gluon")`.
aiter/ops/triton/gluon/gemm_a8w8_blockscale.py	Refactors config caching/copying using shared JSON loader.
aiter/ops/triton/gemm/fused/fused_gemm_afp4wfp4_a16w16.py	Adds device capability gate for FP4 availability.
aiter/ops/triton/gemm/batched/batched_gemm_a8w8_a_per_token_group_prequant_w_per_batched_tensor_quant.py	Stops forcing `kpack` in code (moved to configs).
aiter/ops/triton/gemm/basic/gemm_afp4wfp4.py	Fixes preshuffle path to pass the correct K units to `_get_config`.
aiter/ops/triton/gemm/basic/gemm_a16w16_atomic.py	Removes compatibility defaults for missing config keys (now relies on configs).
aiter/ops/triton/attention/mla_decode_rope.py	Deep-copies shared cached config to avoid mutation of cached dict.
aiter/ops/triton/attention/lean_atten.py	Aligns wrapper config fetch with updated `_get_config()` signature.
aiter/ops/triton/_triton_kernels/attention/mla_decode_rope.py	Switches config load to `load_config_json`.
aiter/ops/triton/_triton_kernels/attention/mha.py	Switches config load to `load_config_json`.
aiter/ops/triton/_triton_kernels/attention/mha_onekernel_bwd.py	Switches config load to `load_config_json`.
aiter/ops/triton/_triton_kernels/attention/mha_fused_bwd.py	Switches config load to `load_config_json`.
aiter/ops/triton/_triton_kernels/attention/lean_atten.py	Uses `load_config_json` and returns a per-call copy to prevent mutation leaks.
aiter/ops/triton/_triton_kernels/attention/hstu_attention.py	Switches config loads to `load_config_json`.
aiter/ops/triton/_triton_kernels/attention/extend_attention.py	Switches config load to `load_config_json`.
aiter/ops/triton/_triton_kernels/moe/moe_routing_sigmoid_top1_fused.py	Switches config load to `load_config_json`.
aiter/ops/triton/_triton_kernels/gmm.py	Switches config load to `load_config_json`.
aiter/ops/triton/_triton_kernels/gemm/fused/fused_gemm_a8w8_blockscale_split_cat.py	Removes LRU decoration from `_get_config` (now relies on shared caching + per-call copies).
aiter/ops/triton/configs/CLAUDE.md	Updates documentation to reflect new caching + resolver behavior.
aiter/ops/triton/configs/gemm/gfx950-BATCHED_GEMM-A8W8-A_PER_TOKEN_GROUP_PREQUANT_W_PER_BATCHED_TENSOR_QUANT.json	Adds `kpack` field.
aiter/ops/triton/configs/gemm/gfx950-BATCHED_GEMM-A8W8-A_PER_TOKEN_GROUP_PREQUANT_W_PER_BATCHED_TENSOR_QUANT-N=8192-K=8192.json	Adds `kpack` field.
aiter/ops/triton/configs/gemm/gfx950-BATCHED_GEMM-A8W8-A_PER_TOKEN_GROUP_PREQUANT_W_PER_BATCHED_TENSOR_QUANT-N=512-K=128.json	Adds `kpack` field.
aiter/ops/triton/configs/gemm/gfx950-BATCHED_GEMM-A8W8-A_PER_TOKEN_GROUP_PREQUANT_W_PER_BATCHED_TENSOR_QUANT-N=128-K=512.json	Adds `kpack` field.
aiter/ops/triton/configs/gemm/gfx942-BATCHED_GEMM-A8W8-A_PER_TOKEN_GROUP_PREQUANT_W_PER_BATCHED_TENSOR_QUANT.json	Adjusts `kpack` from 2 → 1.
aiter/ops/triton/configs/gemm/gfx942-BATCHED_GEMM-A8W8-A_PER_TOKEN_GROUP_PREQUANT_W_PER_BATCHED_TENSOR_QUANT-N=512-K=128.json	Adjusts `kpack` from 2 → 1.
aiter/ops/triton/configs/gemm/gfx942-BATCHED_GEMM-A8W8-A_PER_TOKEN_GROUP_PREQUANT_W_PER_BATCHED_TENSOR_QUANT-N=128-K=512.json	Adjusts `kpack` from 2 → 1.
aiter/ops/triton/configs/gemm/gfx950-GEMM_PREQUANT-AFP4WFP4.json	Removes legacy config file (deleted).
aiter/ops/triton/configs/gemm/gfx950-GEMM_PREQUANT-AFP4WFP4-N=512-K=7168.json	Removes legacy specialized config file (deleted).
aiter/ops/triton/configs/gemm/gfx950-BATCHED_GEMM_PREQUANT-AFP4WFP4.json	Removes legacy config file (deleted).
aiter/ops/triton/configs/gemm/gfx950-BATCHED_GEMM_PREQUANT-AFP4WFP4-N=512-K=128.json	Removes legacy specialized config file (deleted).
aiter/ops/triton/configs/gemm/gfx950-BATCHED_GEMM_PREQUANT-AFP4WFP4-N=128-K=512.json	Removes legacy specialized config file (deleted).
aiter/ops/triton/configs/gemm/gfx1250-GEMM_PREQUANT-AFP4WFP4.json	Removes legacy config file (deleted).
aiter/ops/triton/configs/gemm/gfx1250-BATCHED_GEMM_PREQUANT-AFP4WFP4.json	Removes legacy config file (deleted).

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

brunomazzottiamd · 2026-08-07T19:30:40Z

Added ci:triton-300x label to trigger CI in gfx942. I think it's a good idea since we're dealing with lots of kernels and config files at once.

brunomazzottiamd · 2026-08-07T19:56:21Z

Review in progress...

brunomazzottiamd

Hi @Boss2002n! I'm sharing some comments on you PR.

I've noticed 3 distinct ways of copying config dictionaries, so the caller can mutate them at will without messing with the cached file content:

A. config_dict.copy():

aiter/ops/triton/_triton_kernels/attention/lean_atten.py

B. dict(config_dict): (I think it's equivalent to A.)

aiter/ops/triton/gluon/gemm_a8w8_blockscale.py

C. import copy + copy.deepcopy(config_dict):

aiter/ops/triton/attention/mla_decode_rope.py
op_tests/triton_tests/attention/test_mla_decode_rope.py

I have lots of questions... Can we standardize this config copying? What's the difference between copy and deep copy? Do we really need a deep copy? Is it safer to always rely on deep copy?

Boss2002n · 2026-08-10T14:51:00Z

What's the difference between copy and deep copy? Do we really need a deep copy? Is it safer to always rely on deep copy?

Deep copy creates a new object, shallow copy or just copy references the old object - this is my understanding

brunomazzottiamd

LGTM! Let's get CI green!

brunomazzottiamd · 2026-08-10T15:28:09Z

What's the difference between copy and deep copy? Do we really need a deep copy? Is it safer to always rely on deep copy?

Deep copy creates a new object, shallow copy or just copy references the old object - this is my understanding

So, what's the best option for us? Each case requires a specific method or can we default to doing deep or shallow copy for everything?

brunomazzottiamd · 2026-08-10T15:28:48Z

I think we just need to remove kpack=1 from gfx950 config files and then we're good to merge.

Boss2002n · 2026-08-10T15:32:26Z

What's the difference between copy and deep copy? Do we really need a deep copy? Is it safer to always rely on deep copy?

Deep copy creates a new object, shallow copy or just copy references the old object - this is my understanding

So, what's the best option for us? Each case requires a specific method or can we default to doing deep or shallow copy for everything?

I do think we need to go case by case - I have noticed some places use shallow copy and some use deep copy - I have never looked in depth into it

Boss2002n · 2026-08-10T15:47:43Z

I think we just need to remove kpack=1 from gfx950 config files and then we're good to merge.

#4662

brunomazzottiamd · 2026-08-10T16:59:16Z

What's the difference between copy and deep copy? Do we really need a deep copy? Is it safer to always rely on deep copy?

Deep copy creates a new object, shallow copy or just copy references the old object - this is my understanding

So, what's the best option for us? Each case requires a specific method or can we default to doing deep or shallow copy for everything?

I do think we need to go case by case - I have noticed some places use shallow copy and some use deep copy - I have never looked in depth into it

Fine, let's keep the code as-is for now!

brunomazzottiamd · 2026-08-10T16:59:30Z

I think we just need to remove kpack=1 from gfx950 config files and then we're good to merge.

#4662

Approved!

Boss2002n · 2026-08-10T21:58:10Z

I see that there is an issue with Shard 7 prior to this PR - hence I am going ahead with the merge

Boss2002n force-pushed the satya/config_file_changes branch from 72dc22f to 7d5ce9e Compare August 7, 2026 17:10

Boss2002n marked this pull request as ready for review August 7, 2026 19:04

Boss2002n requested review from a team and a lite review from Copilot August 7, 2026 19:04

Copilot started reviewing on behalf of Boss2002n August 7, 2026 19:06 View session

Copilot AI reviewed Aug 7, 2026

View reviewed changes

brunomazzottiamd added the ci:triton-300x label Aug 7, 2026

zufayu requested a review from vgokhale August 10, 2026 01:16

brunomazzottiamd reviewed Aug 10, 2026

View reviewed changes

Base automatically changed from satya/config_changes to main August 10, 2026 14:34

Boss2002n added 4 commits August 10, 2026 09:34

unify the loading and add a UT for configs

b68463c

formatting

7160dd7

config file and loader changes, remove unused configs

8cc8711

format

6cdf0da

k50112113 force-pushed the satya/config_file_changes branch from 79ae3d9 to 6cdf0da Compare August 10, 2026 14:35

Boss2002n added 2 commits August 10, 2026 14:53

address comments

6009733

address comment

0d828a4

Boss2002n mentioned this pull request Aug 10, 2026

[Bugfix][Triton] Set the _get_config memo only after the config JSON loads #4447

Closed

1 task

brunomazzottiamd approved these changes Aug 10, 2026

View reviewed changes

Boss2002n merged commit 38e534a into main Aug 10, 2026
97 of 104 checks passed

Boss2002n deleted the satya/config_file_changes branch August 10, 2026 21:59

Boss2002n changed the title ~~unify the loading and add a UT for configs~~ [Triton] unify the loading and add a UT for configs Aug 11, 2026

Uh oh!

Conversation

Boss2002n commented Aug 7, 2026

Motivation

Technical Details

Test Plan

Test Result

Submission Checklist

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

brunomazzottiamd commented Aug 7, 2026

Uh oh!

brunomazzottiamd commented Aug 7, 2026

Uh oh!

brunomazzottiamd left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Boss2002n commented Aug 10, 2026

Uh oh!

brunomazzottiamd left a comment

Choose a reason for hiding this comment

Uh oh!

brunomazzottiamd commented Aug 10, 2026

Uh oh!

brunomazzottiamd commented Aug 10, 2026

Uh oh!

Boss2002n commented Aug 10, 2026

Uh oh!

Boss2002n commented Aug 10, 2026

Uh oh!

brunomazzottiamd commented Aug 10, 2026

Uh oh!

brunomazzottiamd commented Aug 10, 2026

Uh oh!

Boss2002n commented Aug 10, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants